use std::mem::size_of;

use nalgebra::Vector2;

use crate::{byte_size, cast_slice, create_shader_from_file};

const BUFFER_CHUNK_SIZE: u64 = 1024;
pub struct PointRender {
    n_points: u32,
    pxs_buffer: wgpu::Buffer,
    render_bind_group: wgpu::BindGroup,
    render_pipeline: wgpu::RenderPipeline,
}

impl PointRender {
    pub fn new(
        camera_buffer: &wgpu::Buffer,
        device: &wgpu::Device,
        format: wgpu::TextureFormat,
        r: f32,
    ) -> Self {
        let render_bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("particle render"),
            entries: &[
                wgpu::BindGroupLayoutEntry {
                    binding: 0,
                    visibility: wgpu::ShaderStages::VERTEX,
                    ty: wgpu::BindingType::Buffer {
                        ty: wgpu::BufferBindingType::Uniform,
                        has_dynamic_offset: false,
                        min_binding_size: None,
                    },
                    count: None,
                }
            ]
        });
        let render_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("particle render"),
            layout: &render_bind_group_layout,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: camera_buffer.as_entire_binding()
                }
            ]
        });
        let render_shader = create_shader_from_file(device, "src/render.wgsl");
        let render_pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: Some("particle render"),
            bind_group_layouts: &[&render_bind_group_layout],
            push_constant_ranges: &[]
        });
        let render_pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
            label: Some("particle render"),
            layout: Some(&render_pipeline_layout),
            vertex: wgpu::VertexState {
                module: &render_shader,
                entry_point: "vs_main",
                compilation_options: wgpu::PipelineCompilationOptions {
                    constants: &[("R".to_owned(), r as f64)].into(),
                    ..Default::default()
                },
                buffers: &[
                    wgpu::VertexBufferLayout {
                        array_stride: size_of::<Vector2<f32>>() as u64,
                        step_mode: wgpu::VertexStepMode::Instance,
                        attributes: &wgpu::vertex_attr_array![0 => Float32x2],
                    }
                ]
            },
            primitive: wgpu::PrimitiveState {
                topology: wgpu::PrimitiveTopology::TriangleStrip,
                ..Default::default()
            },
            depth_stencil: None,
            multisample: Default::default(),
            fragment: Some(wgpu::FragmentState {
                module: &render_shader,
                entry_point: "fs_main",
                compilation_options: Default::default(),
                targets: &[
                    Some(wgpu::ColorTargetState {
                        format,
                        blend: Some(wgpu::BlendState::ALPHA_BLENDING),
                        write_mask: wgpu::ColorWrites::ALL,
                    })
                ]
            }),
            multiview: None,
        });
        let pxs_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("pxs"),
            size: BUFFER_CHUNK_SIZE,
            usage: wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        let n_points = 0;
        Self {
            n_points, pxs_buffer, render_pipeline, render_bind_group
        }
    }
    pub fn update(
        &mut self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        pxs: &[Vector2<f32>]
    ) {
        self.n_points = pxs.len() as u32;
        let pxs_size = byte_size(pxs) as u64;
        if pxs_size > self.pxs_buffer.size() {
            let new_size = pxs_size.div_ceil(BUFFER_CHUNK_SIZE) * BUFFER_CHUNK_SIZE;
            self.pxs_buffer = device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("pxs"),
                usage: wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
                size: new_size,
                mapped_at_creation: true,
            });
            self.pxs_buffer
                .slice(0..pxs_size)
                .get_mapped_range_mut()
                .copy_from_slice(unsafe {
                    cast_slice(pxs)
                });
            self.pxs_buffer.unmap();
        } else {
            queue.write_buffer(&self.pxs_buffer, 0, unsafe {
                cast_slice(pxs)
            });
        }
    }
    pub fn draw<'a>(
        &'a self,
        pass: &mut wgpu::RenderPass<'a>,
    ) {
        pass.set_bind_group(0, &self.render_bind_group, &[]);
        pass.set_pipeline(&self.render_pipeline);
        pass.set_vertex_buffer(0, self.pxs_buffer.slice(..));
        pass.draw(0..4, 0..self.n_points);
    }
}